import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple


class KMeans:
    def __init__(self, n_clusters: int = 2, max_iter: int = 100, tol: float = 1e-4, random_state: int = None):
        """
        初始化KMeans聚类器
        :param n_clusters: 聚类簇数（K值），默认2
        :param max_iter: 最大迭代次数，默认100
        :param tol: 质心变化阈值，小于该值则认为收敛，默认1e-4
        :param random_state: 随机种子，保证结果可复现
        """
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.tol = tol
        self.random_state = random_state
        self.centroids_ = None  # 存储最终质心
        self.labels_ = None  # 存储每个样本的簇标签

    def _euclidean_distance(self, x1: np.ndarray, x2: np.ndarray) -> np.ndarray:
        """
        计算欧氏距离：d = sqrt(sum((x1 - x2)^2))
        :param x1: 形状为(n_samples, n_features)的样本矩阵
        :param x2: 形状为(n_clusters, n_features)的质心矩阵
        :return: 形状为(n_samples, n_clusters)的距离矩阵，每个元素是样本到质心的距离
        """
        # 广播机制：(n_samples,1,n_features) - (1,n_clusters,n_features) → (n_samples,n_clusters,n_features)
        diff = x1[:, np.newaxis, :] - x2[np.newaxis, :, :]
        return np.sqrt(np.sum(diff ** 2, axis=2))

    def _initialize_centroids(self, X: np.ndarray) -> np.ndarray:
        """
        初始化质心：从样本中随机选择K个不重复的点作为初始质心（避免随机生成导致质心偏离数据范围）
        :param X: 输入数据，形状为(n_samples, n_features)
        :return: 初始质心，形状为(n_clusters, n_features)
        """
        np.random.seed(self.random_state)
        n_samples = X.shape[0]
        # 随机选择K个不重复的样本索引
        centroid_indices = np.random.choice(n_samples, size=self.n_clusters, replace=False)
        return X[centroid_indices]

    def fit(self, X: np.ndarray) -> None:
        """
        训练KMeans模型（核心迭代过程）
        :param X: 输入数据，形状为(n_samples, n_features)，n_samples为样本数，n_features为特征数
        """
        # 数据校验：确保输入是2D数组
        if len(X.shape) != 2:
            raise ValueError("输入数据必须是2D数组，形状为(n_samples, n_features)")

        n_samples, n_features = X.shape
        # 初始化质心
        self.centroids_ = self._initialize_centroids(X)

        for _ in range(self.max_iter):
            # 1. 分配簇：计算每个样本到所有质心的距离，取距离最小的簇标签
            distances = self._euclidean_distance(X, self.centroids_)
            self.labels_ = np.argmin(distances, axis=1)  # 形状为(n_samples,)

            # 2. 更新质心：计算每个簇的均值作为新质心
            new_centroids = np.zeros((self.n_clusters, n_features))
            for cluster in range(self.n_clusters):
                # 获取当前簇的所有样本
                cluster_samples = X[self.labels_ == cluster]
                # 若簇中无样本（极端情况），保持原质心
                if len(cluster_samples) == 0:
                    new_centroids[cluster] = self.centroids_[cluster]
                else:
                    new_centroids[cluster] = np.mean(cluster_samples, axis=0)

            # 3. 检查收敛：计算质心变化的最大距离，小于tol则停止迭代
            centroid_shift = np.max(np.sqrt(np.sum((new_centroids - self.centroids_) ** 2, axis=1)))
            if centroid_shift < self.tol:
                break

            # 更新质心
            self.centroids_ = new_centroids

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        对新数据预测簇标签
        :param X: 新数据，形状为(n_samples, n_features)
        :return: 簇标签，形状为(n_samples,)
        """
        if self.centroids_ is None:
            raise RuntimeError("请先调用fit()训练模型")
        distances = self._euclidean_distance(X, self.centroids_)
        return np.argmin(distances, axis=1)


